from torch import nn
import torch


class VGG16_first(nn.Module):
    def __init__(self):
        super(VGG16_first, self).__init__()
        # 第一层保持形状不变，padding需要设为1
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1, stride=1)

    def forward(self, x):
        return self.conv1(x)


if __name__ == "__main__":
    input = torch.randn(1, 3, 224, 224)
    print(input.shape)
    v = VGG16_first()
    output = v(input)
    print(output.shape)
